#include <ros/ros.h>
#include <nav_msgs/Odometry.h>
#include <geometry_msgs/Twist.h>
#include <tf/transform_datatypes.h>
#include <cmath>
#include <vector>
#include <fstream>
#include <string>

// Struct to store data for a single experiment trial
struct ExperimentData {
    double actual_distance; // The target distance for the trial
    double odom_a_x;        // Start odometry position x
    double odom_a_y;        // Start odometry position y
    double odom_b_x;        // End odometry position x
    double odom_b_y;        // End odometry position y
    double odom_calculated; // The distance calculated from odometry readings
};

/**
 * @class OdometryCalibrator
 * @brief A class to automatically perform odometry calibration experiments.
 */
class OdometryCalibrator {
public:
    // Constructor
    OdometryCalibrator(ros::NodeHandle& nh);
    // Main run function
    void run();

private:
    // Callback function
    void odomCallback(const nav_msgs::Odometry::ConstPtr& msg);
    
    // Core functionality
    void runExperiments();
    void analyzeAndSaveResults();
    void moveDistance(double distance);
    void waitForStationary(double timeout = 5.0);
    double calculateDistance(double x1, double y1, double x2, double y2);

    // ROS Handles
    ros::NodeHandle nh_;
    ros::Subscriber odom_sub_;
    ros::Publisher cmd_vel_pub_;
    
    // State variables
    nav_msgs::Odometry current_odom_;
    bool odom_received_ = false;

    // Experiment parameters (loaded from ROS parameter server)
    std::vector<double> target_distances_;
    int num_trials_;
    double linear_speed_;
    std::string odom_topic_;
    std::string cmd_vel_topic_;
    std::string output_file_path_;

    // Experiment results
    std::vector<ExperimentData> experiment_results_;
};

// Constructor implementation
OdometryCalibrator::OdometryCalibrator(ros::NodeHandle& nh) : nh_(nh) {
    // Load parameters from the ROS parameter server, with default values
    ros::NodeHandle private_nh("~"); // Private node handle to get private parameters
    private_nh.param("linear_speed", linear_speed_, 0.2);
    private_nh.param("num_trials", num_trials_, 3);
    private_nh.param("odom_topic", odom_topic_, std::string("/odom"));
    private_nh.param("cmd_vel_topic", cmd_vel_topic_, std::string("/cmd_vel"));
    private_nh.param("output_file_path", output_file_path_, std::string("odometry_calibration_results.csv"));
    
    // Load the list of target distances
    if (!private_nh.getParam("target_distances", target_distances_)) {
        ROS_INFO("Parameter 'target_distances' not found, using default: [0.8, 1.6, 2.4]");
        target_distances_ = {0.8, 1.6, 2.4};
    }

    // Initialize publishers and subscribers
    odom_sub_ = nh_.subscribe(odom_topic_, 10, &OdometryCalibrator::odomCallback, this);
    cmd_vel_pub_ = nh_.advertise<geometry_msgs::Twist>(cmd_vel_topic_, 10);

    ROS_INFO("Odometry Calibrator initialized.");
    ROS_INFO("Linear Speed: %.2f m/s", linear_speed_);
    ROS_INFO("Trials per distance: %d", num_trials_);
    ROS_INFO("Results will be saved to: %s", output_file_path_.c_str());
}

// Odometry callback implementation
void OdometryCalibrator::odomCallback(const nav_msgs::Odometry::ConstPtr& msg) {
    current_odom_ = *msg;
    odom_received_ = true;
}

// Distance calculation implementation
double OdometryCalibrator::calculateDistance(double x1, double y1, double x2, double y2) {
    return sqrt(pow(x2 - x1, 2) + pow(y2 - y1, 2));
}

// waitForStationary implementation
void OdometryCalibrator::waitForStationary(double timeout) {
    ros::Rate rate(10);
    ros::Time start_time = ros::Time::now();
    nav_msgs::Odometry initial_odom;
    
    // Make sure we have the latest odometry data before starting
    ros::spinOnce();
    initial_odom = current_odom_;
    
    ROS_INFO("Waiting for the robot to become stationary...");
    
    while (ros::ok()) {
        if ((ros::Time::now() - start_time).toSec() > timeout) {
            ROS_WARN("Timed out waiting for the robot to be stationary. Proceeding anyway.");
            break;
        }
        
        double dist_moved = calculateDistance(
            initial_odom.pose.pose.position.x, initial_odom.pose.pose.position.y,
            current_odom_.pose.pose.position.x, current_odom_.pose.pose.position.y
        );
        
        if (dist_moved < 0.01) { // Consider stationary if moved less than 1cm in 0.1s
            ROS_INFO("Robot is stationary.");
            break;
        }
        
        initial_odom = current_odom_;
        ros::spinOnce();
        rate.sleep();
    }
}

// moveDistance implementation
void OdometryCalibrator::moveDistance(double distance) {
    if (!odom_received_) {
        ROS_ERROR("Cannot move, no odometry data received.");
        return;
    }
    
    ros::Rate rate(20); // Increase control frequency
    
    // Record starting position
    nav_msgs::Odometry start_odom = current_odom_;
    
    geometry_msgs::Twist cmd_vel;
    cmd_vel.linear.x = linear_speed_;
    cmd_vel.angular.z = 0.0;
    
    ROS_INFO("Commanding robot to move forward %.2f meters...", distance);
    
    double distance_moved = 0.0;
    while (ros::ok() && distance_moved < distance) {
        cmd_vel_pub_.publish(cmd_vel);
        ros::spinOnce();
        
        distance_moved = calculateDistance(
            start_odom.pose.pose.position.x, start_odom.pose.pose.position.y,
            current_odom_.pose.pose.position.x, current_odom_.pose.pose.position.y
        );
        rate.sleep();
    }
    
    // Send stop command
    cmd_vel.linear.x = 0.0;
    cmd_vel_pub_.publish(cmd_vel);
    ros::Duration(0.5).sleep(); // Ensure the stop command is received
    cmd_vel_pub_.publish(cmd_vel);

    ROS_INFO("Move command finished. Odometry shows moved: %.4f meters", distance_moved);
}

// runExperiments implementation
void OdometryCalibrator::runExperiments() {
    for (double target : target_distances_) {
        for (int i = 0; i < num_trials_; ++i) {
            ROS_INFO("\n=======================================================");
            ROS_INFO("Starting Experiment: Target %.2f m, Trial %d/%d", target, i + 1, num_trials_);
            ROS_INFO("=======================================================");

            waitForStationary();
            nav_msgs::Odometry odom_a = current_odom_;
            ROS_INFO("Recording start point (a): (%.4f, %.4f)", odom_a.pose.pose.position.x, odom_a.pose.pose.position.y);
            
            moveDistance(target);
            
            waitForStationary();
            nav_msgs::Odometry odom_b = current_odom_;
            ROS_INFO("Recording end point (b): (%.4f, %.4f)", odom_b.pose.pose.position.x, odom_b.pose.pose.position.y);

            double odom_dist = calculateDistance(
                odom_a.pose.pose.position.x, odom_a.pose.pose.position.y,
                odom_b.pose.pose.position.x, odom_b.pose.pose.position.y
            );
            
            ExperimentData data;
            data.actual_distance = target;
            data.odom_a_x = odom_a.pose.pose.position.x;
            data.odom_a_y = odom_a.pose.pose.position.y;
            data.odom_b_x = odom_b.pose.pose.position.x;
            data.odom_b_y = odom_b.pose.pose.position.y;
            data.odom_calculated = odom_dist;
            experiment_results_.push_back(data);

            ROS_INFO("--- Single Trial Result ---");
            ROS_INFO("Target Distance: %.2f m", target);
            ROS_INFO("Odometry Calculated Distance: %.4f m", odom_dist);
            ROS_INFO("Error: %.4f m", fabs(target - odom_dist));
            
            ROS_INFO("Waiting 5 seconds for the next trial...");
            ros::Duration(5.0).sleep();
        }
    }
}

// analyzeAndSaveResults implementation
void OdometryCalibrator::analyzeAndSaveResults() {
    ROS_INFO("\n\n#################### Experiment Data Analysis ####################");
    
    if (experiment_results_.empty()) {
        ROS_WARN("No experiment data to analyze.");
        return;
    }

    // Simple linear model: actual_distance = k * odom_distance + b
    double sum_odom = 0.0, sum_actual = 0.0, sum_odom_sq = 0.0, sum_odom_actual = 0.0;
    int n = experiment_results_.size();
    
    for (const auto& data : experiment_results_) {
        sum_odom += data.odom_calculated;
        sum_actual += data.actual_distance;
        sum_odom_sq += pow(data.odom_calculated, 2);
        sum_odom_actual += data.odom_calculated * data.actual_distance;
    }
    
    double k = 1.0, b = 0.0;
    double denominator = n * sum_odom_sq - pow(sum_odom, 2);
    if (fabs(denominator) > 1e-6) {
        k = (n * sum_odom_actual - sum_odom * sum_actual) / denominator;
        b = (sum_actual / n) - k * (sum_odom / n);
    } else {
        ROS_WARN("Odometry data did not change, cannot perform linear regression.");
    }
    
    ROS_INFO("Fitted Correction Model: Actual_Distance = %.6f * (Odom_Distance) + %.6f", k, b);
    
    std::ofstream outfile(output_file_path_);
    if (!outfile.is_open()) {
        ROS_ERROR("Failed to open file for writing: %s", output_file_path_.c_str());
        return;
    }

    outfile << "# Odometry Calibration Results" << std::endl;
    outfile << "# Correction Model: Actual_Distance = k * (Odom_Distance) + b" << std::endl;
    outfile << "k," << k << std::endl;
    outfile << "b," << b << std::endl;
    outfile << "# Detailed Data" << std::endl;
    outfile << "Target(m),Odom_Distance(m),Predicted_Distance(m),Prediction_Error(m)" << std::endl;

    double total_error = 0.0;
    for (const auto& data : experiment_results_) {
        double predicted = k * data.odom_calculated + b;
        double error = fabs(predicted - data.actual_distance);
        total_error += error;
        outfile << data.actual_distance << "," << data.odom_calculated << "," << predicted << "," << error << std::endl;
    }
    
    outfile.close();
    ROS_INFO("Average Model Error: %.4f m", total_error / n);
    ROS_INFO("Experiment results successfully saved to: %s", output_file_path_.c_str());
}

// Main run function implementation
void OdometryCalibrator::run() {
    ROS_INFO("Waiting to receive the first odometry message...");
    while (ros::ok() && !odom_received_) {
        ros::spinOnce();
        ros::Duration(0.1).sleep();
    }
    
    if (!ros::ok()) return;
    
    ROS_INFO("Successfully received odometry message. Calibration will start now!");
    ros::Duration(2.0).sleep();

    runExperiments();
    analyzeAndSaveResults();
    
    ROS_INFO("#################### Calibration Program Finished ####################");
}

// Program entry point
int main(int argc, char **argv) {
    ros::init(argc, argv, "odometry_calibrator_node");
    ros::NodeHandle nh;
    
    OdometryCalibrator calibrator(nh);
    calibrator.run();
    
    return 0;
}
